# -*- coding: utf-8 -*-

import json
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from sse_starlette.sse import ServerSentEvent, EventSourceResponse
import uvicorn
from llm_knowledge_base import chat_base,ChangeName
from restful import  RequestData
from dataStorage import message_storage
import os
from starlette.responses import FileResponse as StarletteFileResponse
from datetime import datetime
import logging
from config import cfg



app = FastAPI(
    title='大语言模型后端接口',
    version='0.0.1')
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"])




# url = "http://127.0.0.1:8100/v1/chat/completions"
@app.post(path="/v1/chat/completions",
          summary='对话接口')
async def chat_completion(request: RequestData):
    """
    请求示例
    {
        "model": "gpt-3.5-turbo",
        "messages": [
            {
                "role": "user",
                "content": "请介绍下你自己？"
            }
        ],
        "max_tokens": 1000,
        "temperature": 0.5,
        "stream": true
    }
    """
    def decorate(generator):
        # 流式返回生成内容
        contents = []
        links = []
        scores = []
        for item in generator:
            if isinstance(item, str):#
                contents = [item]
                resp = json.dumps({"id": "0000", "object": "", "created":0000, "model": "",
                                   "choices": [{"index": 0000, "delta": {"role": "assistant", "content": item},
                                                "finish_reason": "err"}]})
                # yield ServerSentEvent(data=resp)
                yield resp
            is_end = item.choices[0].finish_reason
            line = item.choices[0].delta.content
            new_item = {"id": item.id, "object": item.object, "created": item.created, "model": item.model,
                 "choices": [{"index": item.choices[0].index,
                              "delta": {"role": "assistant", "content": line},
                              "finish_reason": is_end}]}
            if is_end == "stop":
                for il in line:
                    score = il.split("**")[-1]
                    link = il.split("**")[0]
                    new_item["choices"][0]["index"] = score
                    scores.append(score)
                    new_item["choices"][0]["delta"]["content"] = link
                    links.append(link)
                    print('返回结果结尾项', new_item)
                    logging.info(f"返回结果结尾项: {new_item}")
                    yield json.dumps(new_item)
                line = ''
            else:
                print(f"中间项",new_item)
                yield json.dumps(new_item)
            contents.append(line)

        print(f"回答的答案是：{"".join(contents)}")
        logging.info(f"回答的答案是：{"".join(contents)}")
        # 对话保存至数据库
        now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        question = request.messages[0].content
        modul = cfg.readValue("map",request.bot)
        message_storage(modul=modul,question=question,additional=json.dumps(links,ensure_ascii=False),
                        answer="".join(contents),user_id=request.chatId)

    print('对话开始')
    request_dict = json.loads(request.json(exclude_unset=True))
    print('请求数据: ', request_dict)
    logging.info('对话开始')
    logging.info(f"请求数据: {request_dict}")
    results = chat_base(request)
    return EventSourceResponse(decorate(results))


@app.post(path="/v2/chat/completions",
          summary='对话接口')
async def chat_completion(request: RequestData):
    def decorate(generator):
        # 流式返回生成内容
        contents = []
        links = []
        scores = []
        i = 1
        for item in generator:
            if isinstance(item, str):#
                contents = [item]
                resp =  [00, item, "err"]
                yield resp #
            is_end = item.choices[0].finish_reason
            line =  item.choices[0].delta.content
            if is_end == "stop":
                for il in line:
                    score = il.split("**")[-1]
                    link = il.split("**")[0]
                    scores.append(score)
                    links.append(link)
                    item = [score, link, "stop"]
                    print('返回结果结尾项', item)
                    logging.info(f"返回结果结尾项: {item}")
                    yield item
                line = ''
            else:
                yield [i ,line, "null"]
            contents.append(line)
            i = i + 1
        print(f"回答的答案是：{"".join(contents)}")
        logging.info(f"回答的答案是：{"".join(contents)}")

        # 对话保存至数据库
        now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        question = request.messages[0].content
        modul = cfg.readValue("map",request.bot)


        message_storage(modul=modul,question=question,additional=json.dumps(links,ensure_ascii=False),
                        answer="".join(contents),user_id=request.chatId)

    print('对话开始')
    request_dict = json.loads(request.json(exclude_unset=True))
    print('请求数据: ', request_dict)
    logging.info('对话开始')
    logging.info(f"请求数据: {request_dict}")
    results = chat_base(request)

    return EventSourceResponse(decorate(results))


@app.get("/download/{filename}")
async def download_file(filename: str):  
    file_path = f"./template/{filename}"  
    return FileResponse(file_path)


@app.get("/view/{filename}")
async def preview_file(filename: str):
    with open(os.path.join(os.getcwd(),"pdf_files",'cte.json'), 'r', encoding='utf-8') as file:
        json_str = file.read()
    data = json.loads(json_str)
    r_f = data[filename]
    print(r_f)
    file_path = os.path.join(os.getcwd(),"pdf_files",r_f)
    print(file_path)
    if os.path.exists(file_path):
        if filename.endswith('.pdf'):
            media_type = 'application/pdf;charset=utf-8'
        elif filename.endswith('.txt'):
            media_type = 'text/plain;charset=utf-8'
        else:
            media_type = 'application/octet-stream'  # 默认二进制流

        return FileResponse(
            path=str(file_path),
            media_type=media_type,
            headers={"Content-Disposition": f"inline; filename={r_f}"}  # 设置为inline以在线预览
        )
    else:
        return {"error": "File not found"}


@app.get("/download_pdf/{pdf_name}")
async def download_file(pdf_name: str):
    pdf_path = os.path.join(os.getcwd(), 'out_search_pdf',pdf_name)
    print(pdf_path)
    if not os.path.exists(pdf_path):
        pdf_path = os.path.join(os.getcwd(), 'gfh_search_pdf', pdf_name)
    if not os.path.exists(pdf_path):
        new_name = ChangeName.sum_len(pdf_name)
        use_name = ChangeName.replace_start(new_name)
        pdf_path = os.path.join(os.getcwd(), 'inner_search_pdf',use_name)
    if os.path.exists(pdf_path):
        print(pdf_path)
        return StarletteFileResponse(path=pdf_path, media_type='application/pdf')
    else:
        return {"error": "File not found"}, 404


@app.get("/download_pdf_v2/{pdf_name}")
async def download_file(pdf_name: str):
    pdf_path = os.path.join(os.getcwd(), 'account_pdf',pdf_name)
    if os.path.exists(pdf_path):
        print(pdf_path)
        return StarletteFileResponse(path=pdf_path, media_type='application/pdf')
    else:
        return {"error": "File not found"}, 404


uvicorn.run(app=app, 
            host='0.0.0.0', 
            port=8100)
